feat(vlm): route VLM GRPO through TQ trainer when data_plane.enabled - #2957

The TQ sync trainer's ``sync_rollout_actor`` silently dropped
PackedTensor multimodal fields (pixel_values, image_grid_thw,
mm_token_type_ids for Qwen2.5-VL / Qwen3-VL, token_type_ids for
Gemma3) because the write path filtered by
``isinstance(v, torch.Tensor)`` — ``PackedTensor`` is not a
``torch.Tensor`` subclass. Consequence: trainer's ``prev_logprobs``
were computed without image embeddings while vLLM's
``generation_logprobs`` used them, producing on Qwen3.5-A3B-Base +
geometry3k:
  * train/token_mult_prob_error avg ~176 (should be ~1.02)
  * train/sampling_importance_ratio ~0.982 uniform (should be ~1.000)
  * train/policy_kl_error avg ~175 (should be ~0.02)
Loss was on-policy (probs_ratio = 1) so training didn't crash, but
the diagnostic metrics flagged the mismatch clearly.

Fix threads multimodal fields end-to-end via a binary wire dispatch
(torch.Tensor | np.ndarray[object]) with these pieces:

* nemo_rl/data_plane/field_registry.py — new FieldSpec (alignment /
  container / stages) collapses five sprinkled schema constants
  (DP_TRAIN_FIELDS, LP_SEED_FIELDS, TOKEN_ALIGNED_FIELDS,
  ADDITIONAL_OPTIONAL_KEY_TENSORS, LP_STAGE_EXCLUDED) into one
  registry. Structural container check catches mis-registrations
  loudly at write time.

* nemo_rl/experience/sync_rollout_actor.py — write-side converts each
  PackedTensor to torch.nested.NestedTensor (jagged layout) +
  companion ``<key>__lengths`` int32 tensor before crossing the wire
  boundary. None entries are padded with zero-length tensors to keep
  the nested batch dim aligned with lengths (mixed text/image
  batches). SmolVLM (dim_to_pack=1) fails loudly with
  NotImplementedError — ragged_idx threading is a follow-up.

* nemo_rl/distributed/batched_data_dict.py::get_multimodal_dict —
  read-side reconstructs the per-sample PackedTensor by slicing the
  materialize-padded rectangular tensor at ``__lengths[i]``.
  Defensive assertion on batch-dim alignment between the parent and
  companion.

* nemo_rl/data_plane/codec.py — pack_jagged_fields dispatch is binary
  (Tensor + ndarray[object]); errors loudly on any other type. No
  packed_tensor container category.

* nemo_rl/data_plane/column_io.py::kv_first_write — filter matches
  the codec's binary universe (Tensor incl. torch.nested; ndarray[object]).

* nemo_rl/models/policy/tq_policy.py::_logprob_dispatch — fetches
  via field_registry.fields_for_stage("logprob" | "ref_lp") instead
  of a hardcoded include list. New multimodal keys drop into the
  registry as one line each; no per-stage constant updates needed.

Diagnostic sub-timers (rollout / flatten / bulk_assemble /
kv_first_write / finish_gen) emit under timing/rollout/sub_* so the
driver-side ``generation`` timer can be decomposed for perf analysis.

Empirical verification on Qwen3.5-A3B-Base + geometry3k (2N x 8G,
container nemo-rl:nightly-07242026):
  * SIR 1.00005 / 1.00002 / 0.99998 (bit-for-bit legacy parity)
  * tmpe step 1..3 = 1.019 / 1.020 / 1.057
  * mean(3-9) total_step_time: TQ 122.0s vs Legacy 125.3s — **TQ is
    3.3s faster at steady state**. sub_kv_first_write dropped from
    ~2.0s (np.ndarray[object] of PackedTensor wrappers, per-object
    pickle) to ~0.64s (native torch.nested storage buffer). Ray RPC
    (driver ↔ actor) ~2.8s, TQ put_samples ~0.64s, net TQ tax on
    generation ~3.4s, offset by ~5s savings in policy_training from
    per-worker TQ fetch. Net TQ throughput ~2.6% faster than legacy.

Known limitations documented in code:
  * SmolVLM ``dim_to_pack=1`` support requires ragged_idx handling
    (follow-up).
  * mm_token_type_ids inclusion in the logprob fetch is what
    actually fixes Qwen2.5-VL / Qwen3-VL 3D RoPE positional encoding
    for image tokens on the trainer side.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

Adds three tests exercising the VLM multimodal wire path added by the
preceding commit:

- `test_get_multimodal_dict_full_vlm_wire_roundtrip` — build a
  realistic batch (pixel_values / image_grid_thw as PackedTensor with
  a None entry, mm_token_type_ids as per-token rectangular tensor,
  plus non-multimodal input_ids/token_mask). Encode via
  `to_nested_wire`, simulate materialize (`to_padded_tensor`), decode
  via `from_nested_wire`. Assert every packed field's `.as_tensor()`
  matches pre-wire, per-token fields pass through unchanged, and
  non-multimodal keys are silently skipped. Covers both
  `as_tensors=True` and `as_tensors=False` paths.

- `test_get_multimodal_dict_missing_companion_asserts` — wire-form
  parent without its `__lengths` companion must raise with a
  wire-contract message, not a bare KeyError deep in trainer forward.

- `test_get_multimodal_dict_empty_batch_skips_wire_field` — 0-row DP
  shard: `from_nested_wire` returns None and the read side skips the
  field (no crash on `PackedTensor.__init__`'s len>0 assert).

Also adds a static registry check: the field names used by the tests
must be in `PACKED_MULTIMODAL_FIELDS` / `PER_TOKEN_MULTIMODAL_FIELDS`
— guards against a future rename that would silently drop the field.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

End-to-end verification through the data-plane ABC contract:
build a rollout batch with PackedTensor pixel_values +
image_grid_thw + mm_token_type_ids → simulate sync_rollout_actor's
write loop → kv_first_write to NoOpDataPlaneClient → read_columns
(materialize) → assert get_multimodal_dict on the fetched batch
matches the pre-wire .as_tensor() output.

Guards the silent-drop regression class this PR was written to fix,
plus the pad_to_seqlen exclusion for multimodal fields (asserts
pixel_values shape stays [B, max_patches, ...] and doesn't inflate
to [B, seqlen, ...]).

Complements the two BatchedDataDict-layer defensive tests in
tests/unit/data/test_multimodal_dict.py (companion-missing and
empty-batch guards).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

Addresses PR #2957 review comments.

1. train_from_meta / train_microbatches_from_meta fetched only the
   static text-only DP_TRAIN_FIELDS, so a VLM training forward ran
   image-blind while prev/ref logprobs were computed *with* images.
   Both now ship DP_TRAIN_FIELDS + _present_multimodal_fields(meta),
   the same per-batch add-on the logprob dispatch uses (renamed
   _LP_MULTIMODAL_FIELDS -> _WIRE_MULTIMODAL_FIELDS since it is no
   longer logprob-only).

   The Qwen3.5-35B-A3B geo3k run shows the fingerprint: train/
   probs_ratio_max 216-2661 and probs_ratio_min 0.0 on every step,
   with probs_ratio_clamped pinned to the 1.2/0.8 clip bounds, while
   the non-DP baseline holds probs_ratio_max == probs_ratio_min ==
   1.0 (the strictly-on-policy invariant for a single inner step).
   Mean probs_ratio stays ~0.998 because only the image tokens are
   wrong, which is why the metrics compared in the PR body -- all
   computed upstream of the training forward -- did not move.

2. BatchedDataDict.truncate_tensors narrowed every >=2-D tensor on
   dim 1 to the microbatch seqlen. The in-memory PackedTensor form is
   skipped by torch.is_tensor, but the data-plane wire form is a
   plain tensor whose dim 1 is patch/image count, so dynamic batching
   would silently corrupt images (or raise narrow: length > size).
   Skip PACKED_MULTIMODAL_FIELDS and their __lengths companions;
   per-token fields stay sequence-aligned and still truncate.

3. Register imgs_sizes (Nemotron-Omni, packed along dim 0 per
   data/processors.py) in PACKED_MULTIMODAL_FIELDS.

Tests: train/logprob multimodal field parity + text-only stays empty
(tests/unit/data_plane/test_multimodal_wire_roundtrip.py), and
truncate_tensors leaving the wire form intact while truncating
mm_token_type_ids (tests/unit/data/test_multimodal_dict.py).

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

run_vlm_grpo passes processor= to whichever trainer _select_trainer
returns, but grpo_train_sync did not accept it, so every VLM run with
data_plane.enabled=true died after full model load with:

    TypeError: grpo_train_sync() got an unexpected keyword argument
    'processor'

grpo_train uses processor for exactly one thing:
attach_initial_nemo_gym_image_payloads, gated on
grpo.deduplicate_multimodal_data. That flag is already rejected for
data_plane.enabled=true by _validate_multimodal_dedup_capability, so the
sync trainer can never need the object -- accept it for signature parity
and assert the invariant instead of ignoring the argument, so relaxing
that upstream guard fails loudly rather than silently dropping image
payloads.

Adds test_sync_trainer_is_call_compatible_with_legacy_trainer, which
diffs the two signatures. The e2e that surfaced this costs two nodes and
~12 minutes of setup before it fails; the signature check is instant.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

The DTensor/AutoModel backend has its own sequence-dim pre-flight,
check_sequence_dim, which is the analog of megatron's
get_and_validate_seqlen. Only the megatron one was taught to skip
multimodal fields, so a TQ VLM run on the automodel backend died at
step 1:

    AssertionError: Dim 1 must be the sequence dim, expected dim 1=2432
    but got shape torch.Size([32, 1, 3])

[32, 1, 3] is image_grid_thw (batch, num_images, t/h/w). Packed
multimodal fields are never sequence-aligned -- dim 1 is
num_images/num_patches -- and their <key>__lengths wire companions are
1-D. In-memory these ride as PackedTensor and are skipped by
torch.is_tensor, but the data-plane wire form is a plain tensor.

The skip goes inside check_sequence_dim rather than onto its existing
skip_keys parameter: all seven call sites need it, and none of them
should have to know the wire format.

Also collapses three byte-identical inline copies of the same check in
the v1 DTensor worker into calls to check_sequence_dim. v1 is the
default worker (dtensor_cfg._v2 defaults to false), so it had the same
latent bug; sharing the helper keeps the skip in one place.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

Multi-agent review of the VLM/TQ path. Each item below is either
reproduced on hardware or verified against source; severities are
post-adversarial-review.

1. mooncake_cpu was dead on arrival for any VLM run. The
   `<key>__lengths` companions minted by `PackedTensor.to_nested_wire`
   are dense int32[B], and `_promote_1d_leaves` rejects any dense 1-D
   field not declared in PROMOTE_1D_FIELDS. Reproduced on 1 node:

     ray::SyncRolloutActor.rollout_to_tq()
     ValueError: Mooncake field 'pixel_values__lengths' is a dense 1D
     tensor but is not declared in data_plane.schema.PROMOTE_1D_FIELDS

   Derived the companion names from PACKED_MULTIMODAL_FIELDS rather
   than hand-listing them, so a new packed modality is covered
   automatically. Fixes the read side too -- `_from_wire` squeezes only
   declared fields. Verified: same job now runs both steps clean.
   Missed because every e2e used backend=simple, where the transform
   is a no-op.

2. The rollout write dropped `pixel_dtype`, which every legacy analog
   passes (grpo.py:2106/:3195, grpo_sync.py:644) and no worker
   re-applies. pixel_values are deliberately fp32 out of the processor
   (see data/processors.py), so they crossed the wire fp32 where legacy
   shipped bf16 -- 2x the largest column. Verified by inspection of all
   call sites, not measured.

3. `to_nested_wire` concatenated a logical row's segments before
   applying pad_to_max_shape, inverting `as_tensor`'s order; a
   multi-segment dedup row with differing trailing dims raised in
   torch.cat before the padding could run. Now pads segments first.
   Latent today (dedup is rejected under data_plane) -- trap removal.

4. `from_nested_wire` turned zero-length rows into empty tensors rather
   than None, so an image-free DP shard produced pixel_values of shape
   (0, ...) where legacy gives None, and logical_segment_counts_by_row
   reported 1 instead of 0. Also added a companion/row length check --
   a short companion would silently misalign images against samples.

5. fp8 QKV calibration filtered on DP_CALIB_INPUT_FIELDS, which names a
   `multi_modal_inputs` column that is never written, so a VLM run
   calibrated image-blind. Newly reachable: run_vlm_grpo only started
   routing to grpo_train_sync in this PR.

6. Docs/typing: the `_logprob_dispatch` docstring stated the opposite of
   what the code does (a reader following it would restore the
   image-blind bug); the registry comment claimed to be the "single
   source of truth" when data/processors.py performs the actual
   classification; a comment cited a test path that does not exist;
   TOKEN_ALIGNED_FIELDS hand-duplicated PER_TOKEN_MULTIMODAL_FIELDS;
   `encode_multimodal_for_wire` was unannotated in a pyrefly-checked
   file.

Not addressed here, both needing a product decision: SmolVLM cannot run
with data_plane.enabled (get_dim_to_pack_along returns 1, which
to_nested_wire rejects) and wants a setup()-time gate; and VLM+TQ has no
nightly coverage because tests/test_suites/llm/common-tq.env gates on
^(grpo|dapo|prorlv2)-, which no vlm_grpo-* recipe can match.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

Closes the coverage gaps the review found. 218 passed on GPU
(tests/unit/data_plane/ + unit/data/test_multimodal_dict.py).

Encoder / wire boundary (test_multimodal_dict.py):
- First direct coverage of `encode_multimodal_for_wire`: parent +
  __lengths emission, per-token passthrough, all-empty skip,
  unregistered-field KeyError, wrong-type guards.
- `to_nested_wire` guard rails: dim_to_pack != 0 -> NotImplementedError,
  all-None -> (None, None), pad_to_max_shape rank mismatch -> ValueError.
- Regression for the pad-before-concat fix: a dedup row spanning 2x4 and
  4x2 segments round-trips instead of raising in torch.cat.
- Regressions for the companion length check and for empty rows matching
  legacy None semantics.
- truncate_tensors leaves wire-form multimodal intact while still
  truncating per-token maps.

Dispatch (test_multimodal_wire_roundtrip.py):
- The roundtrip helper now calls the production
  `encode_multimodal_for_wire` instead of reimplementing its branches --
  the reimplementation was blind by construction to exactly the drift its
  docstring claimed to catch.
- ref_lp and the SC `train_microbatches_from_meta` path each carry their
  own copy of the multimodal add-on and had none; the SC path could have
  regressed image-blind while the sync path stayed green.
- The train/logprob parity assertion now compares against the full
  registry; the previous third assertion was implied by the two above it.
- materialize's pad_to_seqlen exclusion was never actually exercised --
  pad_to_seqlen comes from meta.extra_info, which the test never stamped,
  so the guard against the ~40x blow-up was untested despite the module
  docstring claiming otherwise.
- `_stub_tq_policy` used `__del__ = None`, which does not suppress the
  destructor: CPython installs tp_finalize whenever __del__ is in the
  class dict, then calls None() -> TypeError, surfacing as
  PytestUnraisableExceptionWarning charged to whichever test is at GC.

Architecture invariants (test_architecture_invariants.py):
- run_vlm_grpo's `_select_trainer` copy is now pinned; only run_grpo's was,
  and the VLM launcher is the one that shipped the processor= TypeError.
- The signature-compat check now binds the launcher's actual call shape
  rather than demanding full parameter parity, which would have forced
  every future grpo_train parameter into grpo_train_sync as dead weight.

Seq-dim skip (test_automodel_train.py):
- check_sequence_dim's multimodal skip, plus a negative case proving
  per-token fields are still validated.

The matching get_and_validate_seqlen tests are not included: the
container ships neither megatron.bridge nor transformer_engine, so they
could not be verified here (the pre-existing TestGetAndValidateSeqlen
cannot run in it either). They want the mcore CI shard.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

Two type errors from the Lint check, both introduced by the preceding
review-fix commit:

- `row_segments` typed as list[list[Tensor]] but built by a comprehension
  whose `is not None` filter does not narrow Optional[Tensor], so it
  inferred list[list[Tensor | None]]. Rebuilt with an explicit loop.
- `encode_multimodal_for_wire` yielded the `__lengths` companion while
  only guarding `nested is None`; `to_nested_wire` returns Optional for
  both, so the yield did not match the newly declared
  Iterator[tuple[str, Tensor]]. Guard both.

The second is a direct consequence of annotating the function in that
commit -- the annotation is what exposed the looseness. ruff was clean
throughout, which is why this only surfaced in CI.

Verified with pyrefly in the nightly container (errors shown: 0);
the repo venv here cannot install dev deps.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

…anions

TQ already tracks per-row shape and RL was overriding it in two places.
Both overrides existed to undo something RL itself did, so remove the
cause rather than the compensation.

1. Multimodal `<key>__lengths` companions

`codec.materialize` ran `to_padded_tensor` on every nested leaf,
rectangularizing packed multimodal fields to `[B, max_rows, ...]` and
destroying the row boundaries. The `<key>__lengths` companion field
existed solely to recover them. TQ stores one entry per row
(`storage/managers/base.py::_generate_values` unbinds nested fields), so
the value it hands back already carries the true per-row shapes.

Multimodal fields now skip `to_padded_tensor` and stay nested;
`PackedTensor.from_wire` unbinds instead of slicing by companion length.
That removes a wire field per packed modality, the derived
`PROMOTE_1D_FIELDS` union, the `_WIRE_MULTIMODAL_FIELDS` companion tier,
and the `.endswith(LENGTHS_SUFFIX)` skips in both seq-dim validators and
`truncate_tensors`.

`pad_to_max_shape` is deliberately kept. It is not a transport artifact:
`as_tensor` pads to the same batch max, and a deduplicated logical row can
span segments with differing trailing dims that `torch.cat` cannot join
otherwise.

2. `PROMOTE_1D_FIELDS`

`transfer_queue.metadata.extract_field_schema` rebinds a *local* for 1-D
inputs (`value = value.unsqueeze(-1)`) and derives the sample shape from
it, while `_generate_values` iterates the *original* `(N,)` tensor into N
0-d rows. The schema claims `(1,)`; storage holds `()`. Only the KV path
notices, because it is the only one that reconstructs from the schema --
`SimpleStorage` fetches stored objects and never consults it.

RL compensated by reshaping the payload to match the wrong schema.
`_patch_scalar_field_schema` fixes the schema instead, so the reported
shape matches the stored rows and the column stacks back to a dense
`(N,)`. It rebinds in all three importing modules (both storage managers
bind the name at import time) and is self-verifying: it probes
`KVStorageManager._generate_values` and refuses to install if a TQ
revision ever starts storing 1-D fields as `(1,)` rows.

This covers every dense 1-D field rather than a declared allowlist, so
`PROMOTE_1D_FIELDS`, `_promote_1d_leaves` and `_from_wire`'s squeeze
branch are all deleted.

Tested: 255 passed / 11 skipped across tests/unit/data_plane/ plus
test_multimodal_dict.py; 398 passed on test_batched_data_dict.py,
tests/unit/data/ and test_automodel_train.py. The 11 skips are the
mooncake_cpu fixtures (no RDMA on the test host), so the KV path the
scalar-schema patch targets is NOT yet covered by a run -- that and
Nemotron-Omni's pad_to_max_shape path remain unverified end to end.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Three defects kept VLM GRPO from running on the TransferQueue data plane,
each reproduced before fixing:

* ``BatchedDataDict.slice`` indexed a raw ``torch.nested`` value on its
  ragged dim. A nested tensor satisfies ``isinstance(v, Tensor)`` and has a
  valid ``shape[0]``, so it cleared both guards and failed only at the
  indexing. ``codec.materialize`` now reassembles packed multimodal fields
  into ``PackedTensor`` at the decode boundary, restoring the invariant that
  no raw nested value reaches a consumer.
* ``_from_wire`` densified any nested field whose rows happened to share a
  shape, which for a packed field is a data-dependent accident (every sample
  carrying one image) that discards row boundaries.
* ``_validate_multimodal_dedup_capability`` rejected every
  ``data_plane.enabled=true`` config, blocking all six Nemotron-Omni
  recipes. The gap it guards is NeMo-Gym specific -- ``grpo_train_sync``
  does not call ``attach_initial_nemo_gym_image_payloads``, itself gated on
  ``should_use_nemo_gym`` -- so the check now names that combination.

``to_wire`` additionally flattens each segment to 1-D. Rows then vary only
in dim 0, so ``torch.jagged`` accepts ragged trailing dims and mixed rank
alike; the batch-max padding it used to materialize is gone from the wire
and from TQ storage, and TQ never falls back to the deprecated strided
layout. Padding moves to ``as_tensor``, where a rectangle is actually
required -- it is a model input constraint, not a transport one.

The shapes flattening removes travel on ``KVBatchMeta.tags``, the
transport's per-sample channel, projected with the rows by
``subset``/``slice``/``concat``. ``shard_meta_for_dp`` did not propagate
tags and would have dropped them on every per-rank fetch. Nothing in
``nemo_rl/data_plane`` interprets them.

``_global_pad_shape`` pins the batch-wide pad target so every DP rank sees
identical trailing dims; a rank-local max let the logprob and training
passes encode the same media at different widths. It propagates through all
nine ``PackedTensor`` constructors -- ``concat`` dropping it silently
reverted ``shard_by_batch_size`` to a per-shard max.

Tested:
* Qwen3.5-35B-A3B geo3k 2n8g, 20/20 steps. This recipe trains one inner
  step per rollout, so ``probs_ratio`` is an identity: measured exactly
  1.000000 on every step, across three successive wire formats.
* Nemotron-Omni-30B-A3B clevr 1n8g, 10/10 steps.
  ``token_mult_prob_error`` 1.0138-1.0153 against 1.0138-1.0155 for the
  same recipe with ``data_plane.enabled=false``.
* 245 passed / 11 skipped on tests/unit/data_plane and test_multimodal_dict.

Nightly gates match what each recipe can actually assert: Qwen3.5 gates
``probs_ratio`` exactly; Nemotron-Omni trains 16 inner steps, where that
metric measures policy drift and its max ranges 5.85-29.21 run to run on
identical code (the non-data-plane path included), so it gates
``token_mult_prob_error`` instead.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>

The pad width is scratch the model discards -- mcore crops it via imgs_sizes
before patchification, and the AutoModel path rejects mixed-resolution batches
outright -- so no consumer needs shards to agree on it. Drop the transported
max from the row tags and the _global_pad_shape plumbing that carried it;
as_tensor now computes the max over the rows it holds. Row tags keep shapes
(payload geometry, unrecoverable after to_wire flattens) and pad (the field's
policy flag), so the data plane no longer carries a pad target it cannot
interpret.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>